import gradio as gr
import os
import pathlib
from demo.locales import LOCALES
from demo.processor import IDPhotoProcessor

"""
只裁切模式:
1. 如果重新上传了照片，然后点击按钮，第一次会调用不裁切的模式，第二次会调用裁切的模式
"""


def load_description(fp):
    """
    加载title.md文件作为Demo的顶部栏
    """
    with open(fp, "r", encoding="utf-8") as f:
        content = f.read()
    return content


def create_ui(
    processor: IDPhotoProcessor,
    root_dir: str,
    human_matting_models: list,
    face_detect_models: list,
    language: list,
):

    # 加载环境变量DEFAULT_LANG, 如果有且在language中，则将DEFAULT_LANG设置为环境变量
    if "DEFAULT_LANG" in os.environ and os.environ["DEFAULT_LANG"] in language:
        DEFAULT_LANG = os.environ["DEFAULT_LANG"]
    else:
        DEFAULT_LANG = language[0]

    DEFAULT_HUMAN_MATTING_MODEL = "modnet_photographic_portrait_matting"
    DEFAULT_FACE_DETECT_MODEL = "retinaface-resnet50"

    if DEFAULT_HUMAN_MATTING_MODEL in human_matting_models:
        human_matting_models.remove(DEFAULT_HUMAN_MATTING_MODEL)
        human_matting_models.insert(0, DEFAULT_HUMAN_MATTING_MODEL)

    if DEFAULT_FACE_DETECT_MODEL not in face_detect_models:
        DEFAULT_FACE_DETECT_MODEL = "mtcnn"

    demo = gr.Blocks(title="HivisionIDPhotos")

    with demo:
        gr.HTML(
            "<h1 style='text-align: center; color: #333; margin-bottom: 30px;'>透明高清照生成器</h1>"
        )
        with gr.Row():
            # ------------------------ 左半边 UI ------------------------
            with gr.Column():
                img_input = gr.Image(height=400, label="上传照片")

                with gr.Row():
                    # 简化的人脸检测和抠图模型选择
                    face_detect_model_options = gr.Dropdown(
                        choices=face_detect_models,
                        label="人脸检测模型",
                        value=DEFAULT_FACE_DETECT_MODEL,
                    )

                    matting_model_options = gr.Dropdown(
                        choices=human_matting_models,
                        label="抠图模型",
                        value=human_matting_models[0],
                    )

                # 简化的高级参数
                with gr.Accordion("高级设置", open=False):
                    head_measure_ratio_option = gr.Slider(
                        minimum=0.1,
                        maximum=0.5,
                        value=0.2,
                        step=0.01,
                        label="面部比例",
                        info="调整人脸在照片中的比例",
                    )

                    top_distance_option = gr.Slider(
                        minimum=0.02,
                        maximum=0.5,
                        value=0.12,
                        step=0.01,
                        label="头距顶距离",
                        info="调整头顶到照片边缘的距离",
                    )

            # ---------------- 右半边 UI ----------------
            with gr.Column():
                notification = gr.Text(label="状态信息", visible=False)

                # 只保留透明高清照输出
                with gr.Accordion("透明高清照", open=True):
                    img_output_standard_hd_png = gr.Image(
                        label="透明高清照",
                        height=400,
                        format="png",
                        elem_id="transparent_hd_photo",
                    )

                # 隐藏其他输出（保持兼容性）
                img_output_standard_png = gr.Image(visible=False)
                img_output_standard_hd = gr.Image(visible=False)
                img_output_standard = gr.Image(visible=False)
                img_output_layout = gr.Image(visible=False)
                img_output_template = gr.Image(visible=False)
                template_image_accordion = gr.Accordion(visible=False)
                matting_image_accordion = gr.Accordion(visible=False)

            # ---------------- 简化的处理函数 ----------------
            def process_image(
                input_image,
                matting_model,
                face_detect_model,
                head_ratio,
                top_distance,
            ):
                """简化的图像处理函数，专注透明高清照生成"""
                try:
                    # 调用原有的processor，使用默认参数生成透明高清照
                    result = processor.process(
                        input_image,
                        "只换底",  # mode_option - 只换底模式，不需要尺寸和颜色
                        "一寸照",  # size_list_option (会被忽略)
                        "白色",  # color_option (会被忽略)
                        "纯色",  # render_option
                        "不设置",  # image_kb_options
                        0,  # custom_color_R
                        0,  # custom_color_G
                        0,  # custom_color_B
                        "000000",  # custom_color_hex_value
                        0,  # custom_size_height_px (不使用)
                        0,  # custom_size_width_px (不使用)
                        0,  # custom_size_height_mm (不使用)
                        0,  # custom_size_width_mm (不使用)
                        50,  # custom_image_kb_size
                        "zh",  # language
                        matting_model,
                        "不添加",  # watermark_option
                        "",  # watermark_text
                        "#FFFFFF",  # watermark_text_color
                        20,  # watermark_text_size
                        0.15,  # watermark_text_opacity
                        30,  # watermark_text_angle
                        25,  # watermark_text_space
                        face_detect_model,
                        head_ratio,
                        top_distance,
                        0,  # whitening_strength
                        "不设置",  # image_dpi_options
                        300,  # custom_image_dpi_size
                        0,  # brightness_strength
                        0,  # contrast_strength
                        0,  # sharpen_strength
                        0,  # saturation_strength
                        "六寸",  # print_options
                    )

                    # 返回结果，只保留透明高清照
                    return [
                        result[3],  # img_output_standard_hd_png - 透明高清照
                        result[7],  # notification
                    ]
                except Exception as e:
                    return [None, f"处理错误: {str(e)}"]

            # 绑定图片上传事件，自动生成透明高清照
            img_input.upload(
                process_image,
                inputs=[
                    img_input,
                    matting_model_options,
                    face_detect_model_options,
                    head_measure_ratio_option,
                    top_distance_option,
                ],
                outputs=[
                    img_output_standard_hd_png,
                    notification,
                ],
            )

            # 绑定模型参数变化事件，自动重新生成
            matting_model_options.change(
                process_image,
                inputs=[
                    img_input,
                    matting_model_options,
                    face_detect_model_options,
                    head_measure_ratio_option,
                    top_distance_option,
                ],
                outputs=[
                    img_output_standard_hd_png,
                    notification,
                ],
            )

            face_detect_model_options.change(
                process_image,
                inputs=[
                    img_input,
                    matting_model_options,
                    face_detect_model_options,
                    head_measure_ratio_option,
                    top_distance_option,
                ],
                outputs=[
                    img_output_standard_hd_png,
                    notification,
                ],
            )

            head_measure_ratio_option.release(
                process_image,
                inputs=[
                    img_input,
                    matting_model_options,
                    face_detect_model_options,
                    head_measure_ratio_option,
                    top_distance_option,
                ],
                outputs=[
                    img_output_standard_hd_png,
                    notification,
                ],
            )

            top_distance_option.release(
                process_image,
                inputs=[
                    img_input,
                    matting_model_options,
                    face_detect_model_options,
                    head_measure_ratio_option,
                    top_distance_option,
                ],
                outputs=[
                    img_output_standard_hd_png,
                    notification,
                ],
            )

    return demo
